
#include "DriverMain.h"


#define  DELAY_ONE_MICROSECOND  (-10)
#define  DELAY_ONE_MILLISECOND (DELAY_ONE_MICROSECOND*1000)
#define  DELAY_ONE_SECOND (DELAY_ONE_MILLISECOND*1000)

PDEVICE_OBJECT g_pDeviceObj = NULL;

PIRP  g_pIrpReadPending = NULL;

char HexBuf[] = "0123456789ABCDEF";

void BufferToHex(char* destBuf, char * buffer, int bufferLen) 
{
	if (bufferLen <= 0)
		return;

	RtlZeroMemory(destBuf, bufferLen * 2 + 1);

	char* p = destBuf;
	for (int i =0; i < bufferLen; ++i)
	{
		unsigned char bt = (unsigned char) (buffer[i]);
		*p++ = HexBuf[bt >> 4];
		*p++ = HexBuf[bt & 0x0f];
	}
}


typedef struct _MyDeviceExtension 
{
	PDEVICE_OBJECT attachedDeviceObj;
	UNICODE_STRING attachedDeviceName;
	WCHAR buffer[64];
} MyDeviceExtension, * PMyDeviceExtension;

VOID DriverUnload(IN PDRIVER_OBJECT pDriverObject);

NTSTATUS DeviceDispathFunction(IN PDEVICE_OBJECT pDevice,
							   IN PIRP pIrp);

NTSTATUS AttachedToSerial(IN PDEVICE_OBJECT pDevice, IN int SericalNo);

extern "C"
NTSTATUS DriverEntry(IN PDRIVER_OBJECT pDriverObject, 
					 IN PUNICODE_STRING pRegistryPath)
{
	UNICODE_STRING deviceName = RTL_CONSTANT_STRING(DEVICE_NAME);
	NTSTATUS status = IoCreateDevice(pDriverObject, sizeof(MyDeviceExtension), &deviceName, 
		FILE_DEVICE_UNKNOWN, 0, TRUE, &g_pDeviceObj);
	if (NT_SUCCESS(status))
		KdPrint(("Create Device Success"));
	else
		KdPrint(("Create Device Failure"));
	
// 	UNICODE_STRING symName = RTL_CONSTANT_STRING(SYM_DEVICE_NAME);
// 	status = IoCreateSymbolicLink(&symName, &deviceName);
// 	if (NT_SUCCESS(status))
// 		KdPrint(("IoCreateSymbolicLink Success"));
// 	else
// 		KdPrint(("IoCreateSymbolicLink Failure"));

	if (AttachedToSerial(g_pDeviceObj, 0) != STATUS_SUCCESS)
	{
		//IoDeleteSymbolicLink(&symName);
		KdPrint(("Attach Device Failure"));
		IoDeleteDevice(g_pDeviceObj);
		g_pDeviceObj = NULL;
	}
	else
	{
		KdPrint(("Attach Device Success"));
	}

	for (int i = 0; i < IRP_MJ_MAXIMUM_FUNCTION; ++i)
	{
		pDriverObject->MajorFunction[i] = DeviceDispathFunction;
	}
	


	
	pDriverObject->DriverUnload = DriverUnload;

	return STATUS_SUCCESS;
}


VOID DriverUnload(IN PDRIVER_OBJECT pDriverObject)
{
	//UNICODE_STRING symName = RTL_CONSTANT_STRING(SYM_DEVICE_NAME);
	//IoDeleteSymbolicLink(&symName);

	if (g_pDeviceObj)
	{
		PMyDeviceExtension pExten = (PMyDeviceExtension) g_pDeviceObj->DeviceExtension;

		IoDetachDevice(pExten->attachedDeviceObj);

		LARGE_INTEGER interval;
		interval.QuadPart = (5*1000 * DELAY_ONE_MILLISECOND);        
		KeDelayExecutionThread(KernelMode,FALSE,&interval);
		if (g_pIrpReadPending != NULL)
		{
			//IoCancelIrp(g_pIrpReadPending);
			PIO_STACK_LOCATION irpSp = IoGetNextIrpStackLocation(g_pIrpReadPending);
			irpSp->CompletionRoutine = NULL;
			IoSkipCurrentIrpStackLocation(g_pIrpReadPending);
			g_pIrpReadPending = NULL;
		}

		IoDeleteDevice(g_pDeviceObj);
	}

	KdPrint(("DriverUnload"));
	
}

NTSTATUS ReadCompletionRoutine(IN PDEVICE_OBJECT DeviceObject, IN PIRP Irp, IN PVOID Context)
{
	KdPrint(("ReadCompletionRoutine"));
	g_pIrpReadPending = NULL;
	if (Irp->PendingReturned)
		IoMarkIrpPending(Irp);

	if (NT_SUCCESS(Irp->IoStatus.Status))
	{
		PIO_STACK_LOCATION stack = IoGetCurrentIrpStackLocation(Irp);
		ULONG len = Irp->IoStatus.Information;
		char* buf = NULL;
		//if (Irp->MdlAddress != NULL)
		//	buf = (char*) MmGetSystemAddressForMdlSafe(Irp->MdlAddress, NormalPagePriority);
		//else if (Irp->UserBuffer != NULL)
		//	buf = (char*)Irp->UserBuffer;
		//else if (Irp->AssociatedIrp.SystemBuffer != NULL)
		//	buf = (char*) Irp->AssociatedIrp.SystemBuffer;

		buf = (char*) Irp->AssociatedIrp.SystemBuffer;
		if (buf != NULL)
			KdPrint(("buf != NULL"));
		else
			KdPrint(("buf == NULL"));
		KdPrint(("Read len = %d", len));

		char hex[1025] = {};
		BufferToHex(hex, buf, len);  // todo: check bounds
		KdPrint((hex));
	}
	return Irp->IoStatus.Status;
}

NTSTATUS DeviceDispathFunction(IN PDEVICE_OBJECT pDevice,
							   IN PIRP pIrp)
{
	PIO_STACK_LOCATION stack = IoGetCurrentIrpStackLocation(pIrp);
	//KdPrint(("DeviceDispathFunction"));
	//if (stack->MajorFunction != IRP_MJ_INTERNAL_DEVICE_CONTROL)
	//	KdPrint(("MajorFunction = %S", MJ_FUNC_NAME[stack->MajorFunction]));
	PMyDeviceExtension pExten = (PMyDeviceExtension) pDevice->DeviceExtension;
	PDEVICE_OBJECT pTopObj = pExten->attachedDeviceObj;

	switch (stack->MajorFunction)
	{
	case IRP_MJ_CREATE:
	case IRP_MJ_CLOSE:
		KdPrint(("MajorFunction = %S", MJ_FUNC_NAME[stack->MajorFunction]));

		break;

	case IRP_MJ_POWER:
		PoStartNextPowerIrp(pIrp);
		IoSkipCurrentIrpStackLocation(pIrp);
		return PoCallDriver(pTopObj, pIrp);

	case IRP_MJ_WRITE:
		{
			KdPrint(("MajorFunction = %S", MJ_FUNC_NAME[stack->MajorFunction]));
			ULONG len = stack->Parameters.Write.Length;
			PCHAR buf = NULL;
			if (pIrp->MdlAddress != NULL)
				buf = (PCHAR) MmGetSystemAddressForMdlSafe(pIrp->MdlAddress, NormalPagePriority);
			else if (pIrp->UserBuffer != NULL)
				buf = (PCHAR)pIrp->UserBuffer;
			else if (pIrp->AssociatedIrp.SystemBuffer != NULL)
				buf = (PCHAR) pIrp->AssociatedIrp.SystemBuffer;

			KdPrint(("Write len = %d", len));
			char hex[1025] = {};
			BufferToHex(hex, buf, len);  // todo: check bounds
			KdPrint((hex));
		}
			 
		break;

	case IRP_MJ_READ:
		{
			KdPrint(("MajorFunction = %S", MJ_FUNC_NAME[stack->MajorFunction]));
			IoCopyCurrentIrpStackLocationToNext(pIrp);
			g_pIrpReadPending = pIrp;
			IoSetCompletionRoutine(pIrp, ReadCompletionRoutine, NULL, TRUE, TRUE, TRUE);
			NTSTATUS status = IoCallDriver(pTopObj, pIrp);
			return status;
		}
		break;
	}

	IoSkipCurrentIrpStackLocation(pIrp);
	NTSTATUS status = IoCallDriver(pTopObj, pIrp);

	return status;
}

NTSTATUS AttachedToSerial(IN PDEVICE_OBJECT pDevice, IN int SericalNo)
{
	WCHAR buffer[64] = {};
	UNICODE_STRING objName;
	RtlInitEmptyUnicodeString(&objName, buffer, 256);
	RtlUnicodeStringPrintf(&objName, L"\\Device\\Serial%d", SericalNo);

	PFILE_OBJECT fileObj =  NULL;
	PDEVICE_OBJECT devObj = NULL;
	NTSTATUS status = IoGetDeviceObjectPointer(&objName, FILE_ALL_ACCESS, &fileObj, &devObj);
	if (status != STATUS_SUCCESS)
		return STATUS_UNSUCCESSFUL;
	
	ObDereferenceObject(fileObj);
	if (devObj->Flags & DO_BUFFERED_IO)
		pDevice->Flags |= DO_BUFFERED_IO;

	if (devObj->Flags & DO_DIRECT_IO)
		pDevice->Flags |= DO_DIRECT_IO;

	if (devObj->Characteristics & FILE_DEVICE_SECURE_OPEN)
		pDevice->Characteristics |= FILE_DEVICE_SECURE_OPEN;

	pDevice->Flags |= DO_POWER_PAGABLE;

	PDEVICE_OBJECT topDev = IoAttachDeviceToDeviceStack(pDevice, devObj);
	if (topDev == NULL) 
	{
		return STATUS_UNSUCCESSFUL;
	}

	PMyDeviceExtension pExten = (PMyDeviceExtension) pDevice->DeviceExtension;
	pExten->attachedDeviceObj = topDev;
	RtlInitEmptyUnicodeString(&pExten->attachedDeviceName, &pExten->buffer[0], 64);
	RtlCopyUnicodeString(&pExten->attachedDeviceName, &objName);
	pDevice->Flags = pDevice->Flags & ~ (DO_DEVICE_INITIALIZING);

	// Test
	KdPrint(("g_pDeviceObj = %8X", g_pDeviceObj));
	KdPrint(("g_pDeviceObj.AttachedDevice = %8X", g_pDeviceObj->AttachedDevice));
	KdPrint(("TopDev = %8X", topDev));
	KdPrint(("TopDev.AttachedDevice = %8X", topDev->AttachedDevice));

	return STATUS_SUCCESS;
}